# DQN网络模型,包括经验回访池的实现

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from collections import namedtuple,deque
from random import choice, randrange, sample

from config import *

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
Transition = namedtuple('Transition', ('state', 'action', 'state_next', 'reward'))  # Define a transition tuple
class DNN(nn.Module):
    def __init__(self, n_state=STATE_DIM,
                 n_action=ACTION_DIM):  # Define the layers of the fully-connected hidden network
        super(DNN, self).__init__()
        self.input_layer = nn.Linear(n_state, 1000)
        self.input_layer.weight.data.normal_(0, 0.1)  # 参数初始化
        self.middle_layer = nn.Linear(1000, 500)
        self.middle_layer.weight.data.normal_(0, 0.1)
        self.middle_layer_2 = nn.Linear(500, 250)
        self.adv_layer = nn.Linear(250, n_action)
        self.adv_layer.weight.data.normal_(0, 0.1)

    def forward(self, state):
        x = F.relu(self.input_layer(state))
        x = F.relu(self.middle_layer(x))
        x = F.relu(self.middle_layer_2(x))
        out = self.adv_layer(x)
        return out

class DQN(object):
    def __init__(self):
        self.policy_net = self.creat_model().to(device)
        self.target_net = self.creat_model().to(device)
        self.replay_size = MEMORY_CAPACITY
        self.replay_queue = deque(maxlen=self.replay_size)
        self.learn_step_counter = 0
        self.loss_func = nn.MSELoss()
        self.optimizer = torch.optim.Adam(self.policy_net.parameters(), lr=LR)
    def creat_model(self):
        print("creating DNN.....")
        return DNN()

    def choose_action(self, s, epsilon):
        # 选择动作根据e-greedy策略
        if np.random.uniform() < epsilon:
            return np.random.choice(ACTION_DIM)
        else:
            s = torch.FloatTensor(s).unsqueeze(0)  # 添加额外的batch维度
            if next(self.policy_net.parameters()).is_cuda:
                s = s.cuda()  # 如果模型在GPU上，确保状态也在GPU上

            with torch.no_grad():
                q_values = self.policy_net(s)
                action = torch.argmax(q_values).item()  # 选择价值最高的动作
            return action

    def remember(self, s, a, next_s, reward):
        # save MDP transitions
        self.replay_queue.append(Transition(s, a, next_s, reward))

    def train(self):
        if len(self.replay_queue) < MEMORY_CAPACITY:
            print(f"经验回放池不够大，当前长度：{len(self.replay_queue)}")
            return

        self.learn_step_counter += 1
        transitions = self.sample(BENCH_SIZE)
        batch = Transition(*zip(*transitions))

        # 转换为 PyTorch 张量并移到适当的设备上
        state_batch = torch.tensor(batch.state, dtype=torch.float32).to(device)
        action_batch = torch.tensor(batch.action, dtype=torch.long).unsqueeze(-1).to(device)
        reward_batch = torch.tensor(batch.reward, dtype=torch.float32).unsqueeze(-1).to(device)
        state_next_batch = torch.tensor(batch.state_next, dtype=torch.float32).to(device)

        # 计算当前状态的 Q 值
        state_action_values = self.policy_net(state_batch).gather(1, action_batch)

        # 使用目标网络计算下一个状态的 V 值
        next_state_values = self.target_net(state_next_batch).max(1)[0].detach()

        # 计算预期的 Q 值
        expected_state_action_values = reward_batch + (GAMMA * next_state_values).unsqueeze(1)

        # 计算损失
        loss = F.mse_loss(state_action_values, expected_state_action_values)

        # 优化模型
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        if self.learn_step_counter % PRINT_EVERY == 0:
            print(f"训练迭代 {self.learn_step_counter}，损失：{loss.item()}")

        # 更新目标网络
        if self.learn_step_counter % TARGET_UPDATE == 0:
            self.target_net.load_state_dict(self.policy_net.state_dict())

    def sample(self, batch_size):
        return sample(self.replay_queue, batch_size)


